import os
import time
import numpy as np
from PIL import Image
import io
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import padding

class ImageCrypto:
    """
    图像加密处理类
    支持对图像进行加密和解密操作
    """
    
    def __init__(self, key_size=32):
        self.key_size = key_size
        self.iv_size = 16
    
    def encrypt_image(self, image_path, key, output_path=None):
        """
        加密图像文件
        :param image_path: 输入图像路径
        :param key: 加密密钥
        :param output_path: 输出文件路径，默认为None（返回加密后的字节数据）
        :return: 加密后的数据或加密性能指标
        """
        # 读取图像
        img = Image.open(image_path)
        
        # 转换为字节数据
        img_byte_arr = io.BytesIO()
        img.save(img_byte_arr, format=img.format or 'PNG')
        img_data = img_byte_arr.getvalue()
        
        # 记录开始时间
        start_time = time.time()
        
        # 生成随机IV
        iv = os.urandom(self.iv_size)
        
        # 使用PKCS7填充
        padder = padding.PKCS7(128).padder()
        padded_data = padder.update(img_data) + padder.finalize()
        
        # 使用AES-256-CBC加密
        cipher = Cipher(algorithms.AES(key), modes.CBC(iv))
        encryptor = cipher.encryptor()
        encrypted_data = encryptor.update(padded_data) + encryptor.finalize()
        
        # 计算执行时间和性能
        execution_time = time.time() - start_time
        data_size = len(img_data) / (1024 * 1024)  # 转换为MB
        # 避免除以零错误
        performance = data_size / max(execution_time, 0.001)  # MB/s，确保执行时间至少为1毫秒
        
        # 将IV附加到加密数据前面
        final_data = iv + encrypted_data
        
        if output_path:
            with open(output_path, 'wb') as f:
                f.write(final_data)
            return {
                'execution_time_ms': execution_time * 1000,
                'data_size_mb': data_size,
                'performance_mbps': performance,
                'output_path': output_path
            }
        else:
            return final_data
    
    def decrypt_image(self, encrypted_path, key, output_path=None):
        """
        解密图像文件
        :param encrypted_path: 加密图像路径
        :param key: 解密密钥
        :param output_path: 输出文件路径，默认为None（返回解密后的Image对象）
        :return: 解密后的图像或解密性能指标
        """
        # 读取加密数据
        with open(encrypted_path, 'rb') as f:
            data = f.read()
        
        # 提取IV和加密数据
        iv = data[:self.iv_size]
        encrypted_data = data[self.iv_size:]
        
        # 记录开始时间
        start_time = time.time()
        
        # 使用AES-256-CBC解密
        cipher = Cipher(algorithms.AES(key), modes.CBC(iv))
        decryptor = cipher.decryptor()
        padded_data = decryptor.update(encrypted_data) + decryptor.finalize()
        
        # 去除PKCS7填充
        unpadder = padding.PKCS7(128).unpadder()
        original_data = unpadder.update(padded_data) + unpadder.finalize()
        
        # 计算执行时间和性能
        execution_time = time.time() - start_time
        data_size = len(encrypted_data) / (1024 * 1024)  # 转换为MB
        # 避免除以零错误
        performance = data_size / max(execution_time, 0.001)  # MB/s，确保执行时间至少为1毫秒
        
        # 从字节数据恢复图像
        img = Image.open(io.BytesIO(original_data))
        
        if output_path:
            img.save(output_path)
            return {
                'execution_time_ms': execution_time * 1000,
                'data_size_mb': data_size,
                'performance_mbps': performance,
                'output_path': output_path
            }
        else:
            return img
    
    def encrypt_image_chaotic(self, image_path, key, output_path=None):
        """
        使用混沌加密图像
        :param image_path: 输入图像路径
        :param key: 加密密钥
        :param output_path: 输出文件路径，默认为None（返回加密后的字节数据）
        :return: 加密后的数据或加密性能指标
        """
        # 读取图像为NumPy数组
        img = np.array(Image.open(image_path))
        height, width = img.shape[0], img.shape[1]
        channels = 1 if len(img.shape) == 2 else img.shape[2]
        
        # 记录开始时间
        start_time = time.time()
        
        # 从密钥生成混沌系统的初始参数
        # 使用密钥的前16字节作为种子
        seed = int.from_bytes(key[:16], byteorder='big')
        np.random.seed(seed)
        
        # 生成混沌序列
        chaotic_seq = np.random.randint(0, 256, size=height * width * channels, dtype=np.uint8)
        chaotic_map = chaotic_seq.reshape(img.shape)
        
        # 应用异或操作进行加密
        encrypted_img = np.bitwise_xor(img, chaotic_map)
        
        # 计算执行时间和性能
        execution_time = time.time() - start_time
        data_size = (height * width * channels) / (1024 * 1024)  # 转换为MB
        # 避免除以零错误
        performance = data_size / max(execution_time, 0.001)  # MB/s，确保执行时间至少为1毫秒
        
        # 将NumPy数组转换回Image
        encrypted_image = Image.fromarray(encrypted_img.astype(np.uint8))
        
        if output_path:
            # 额外保存密钥相关信息以便解密
            metadata = {
                'height': height,
                'width': width,
                'channels': channels
            }
            # 保存加密图像
            encrypted_image.save(output_path)
            
            # 创建元数据文件
            metadata_path = f"{os.path.splitext(output_path)[0]}.meta"
            with open(metadata_path, 'wb') as f:
                # 简单加密存储元数据
                metadata_str = f"{height},{width},{channels}"
                iv = os.urandom(self.iv_size)
                cipher = Cipher(algorithms.AES(key), modes.CBC(iv))
                encryptor = cipher.encryptor()
                padder = padding.PKCS7(128).padder()
                padded_data = padder.update(metadata_str.encode()) + padder.finalize()
                encrypted_metadata = encryptor.update(padded_data) + encryptor.finalize()
                f.write(iv + encrypted_metadata)
            
            return {
                'execution_time_ms': execution_time * 1000,
                'data_size_mb': data_size,
                'performance_mbps': performance,
                'output_path': output_path,
                'metadata_path': metadata_path
            }
        else:
            # 将加密图像转换为字节数据
            img_byte_arr = io.BytesIO()
            encrypted_image.save(img_byte_arr, format='PNG')
            return img_byte_arr.getvalue()
    
    def decrypt_image_chaotic(self, encrypted_path, key, metadata_path=None, output_path=None):
        """
        使用混沌解密图像
        :param encrypted_path: 加密图像路径
        :param key: 解密密钥
        :param metadata_path: 元数据文件路径，如果为None，则假设与加密文件同名但扩展名为.meta
        :param output_path: 输出文件路径，默认为None（返回解密后的Image对象）
        :return: 解密后的图像或解密性能指标
        """
        # 读取加密图像
        encrypted_img = np.array(Image.open(encrypted_path))
        
        # 如果未提供元数据路径，则假设与加密文件同名但扩展名为.meta
        if metadata_path is None:
            metadata_path = f"{os.path.splitext(encrypted_path)[0]}.meta"
        
        # 读取并解密元数据
        with open(metadata_path, 'rb') as f:
            metadata_data = f.read()
        
        iv = metadata_data[:self.iv_size]
        encrypted_metadata = metadata_data[self.iv_size:]
        
        cipher = Cipher(algorithms.AES(key), modes.CBC(iv))
        decryptor = cipher.decryptor()
        padded_data = decryptor.update(encrypted_metadata) + decryptor.finalize()
        
        unpadder = padding.PKCS7(128).unpadder()
        metadata_str = unpadder.update(padded_data) + unpadder.finalize()
        
        # 解析元数据
        height, width, channels = map(int, metadata_str.decode().split(','))
        
        # 记录开始时间
        start_time = time.time()
        
        # 从密钥生成相同的混沌序列
        seed = int.from_bytes(key[:16], byteorder='big')
        np.random.seed(seed)
        
        # 生成与加密时相同的混沌序列
        chaotic_seq = np.random.randint(0, 256, size=height * width * channels, dtype=np.uint8)
        chaotic_map = chaotic_seq.reshape(encrypted_img.shape)
        
        # 应用异或操作进行解密
        decrypted_img = np.bitwise_xor(encrypted_img, chaotic_map)
        
        # 计算执行时间和性能
        execution_time = time.time() - start_time
        data_size = (height * width * channels) / (1024 * 1024)  # 转换为MB
        # 避免除以零错误
        performance = data_size / max(execution_time, 0.001)  # MB/s，确保执行时间至少为1毫秒
        
        # 将NumPy数组转换回Image
        decrypted_image = Image.fromarray(decrypted_img.astype(np.uint8))
        
        if output_path:
            decrypted_image.save(output_path)
            return {
                'execution_time_ms': execution_time * 1000,
                'data_size_mb': data_size,
                'performance_mbps': performance,
                'output_path': output_path
            }
        else:
            return decrypted_image 